/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.runners.gaia;



import com.bff.gaia.unified.model.pipeline.v1.RunnerApi;

import com.bff.gaia.unified.runners.core.construction.PipelineOptionsTranslation;

import com.bff.gaia.unified.runners.fnexecution.jobsubmission.JobInvocation;

import com.bff.gaia.unified.runners.fnexecution.jobsubmission.JobInvoker;

import com.bff.gaia.unified.runners.fnexecution.provisioning.JobInfo;

import com.bff.gaia.unified.sdk.options.PortablePipelineOptions;

import com.bff.gaia.unified.vendor.grpc.v1p13p1.com.google.protobuf.Struct;

import com.bff.gaia.unified.vendor.guava.com.google.common.util.concurrent.ListeningExecutorService;

import org.slf4j.Logger;

import org.slf4j.LoggerFactory;



import javax.annotation.Nullable;

import java.io.IOException;

import java.util.List;

import java.util.UUID;



import static com.bff.gaia.unified.runners.core.construction.PipelineResources.detectClassPathResourcesToStage;



/** Job Invoker for the {@link GaiaRunner}. */

public class GaiaJobInvoker extends JobInvoker {

  private static final Logger LOG = LoggerFactory.getLogger(GaiaJobInvoker.class);



  public static GaiaJobInvoker create(GaiaJobServerDriver.GaiaServerConfiguration serverConfig) {

    return new GaiaJobInvoker(serverConfig);

  }



  private final GaiaJobServerDriver.GaiaServerConfiguration serverConfig;



  private GaiaJobInvoker(GaiaJobServerDriver.GaiaServerConfiguration serverConfig) {

    super("gaia-runner-job-invoker");

    this.serverConfig = serverConfig;

  }



  @Override

  protected JobInvocation invokeWithExecutor(

      RunnerApi.Pipeline pipeline,

      Struct options,

      @Nullable String retrievalToken,

      ListeningExecutorService executorService)

      throws IOException {

    // TODO: How to make Java/Python agree on names of keys and their values?

    LOG.trace("Parsing pipeline options");

    GaiaPipelineOptions gaiaOptions =

        PipelineOptionsTranslation.fromProto(options).as(GaiaPipelineOptions.class);



    String invocationId =

        String.format("%s_%s", gaiaOptions.getJobName(), UUID.randomUUID().toString());

    LOG.info("Invoking job {}", invocationId);



    if (GaiaPipelineOptions.AUTO.equals(gaiaOptions.getGaiaMaster())) {

      gaiaOptions.setGaiaMaster(serverConfig.getGaiaMasterUrl());

    }



    PortablePipelineOptions portableOptions = gaiaOptions.as(PortablePipelineOptions.class);

    if (portableOptions.getSdkWorkerParallelism() == 0L) {

      portableOptions.setSdkWorkerParallelism(serverConfig.getSdkWorkerParallelism());

    }



    gaiaOptions.setRunner(null);



    return createJobInvocation(

        invocationId,

        retrievalToken,

        executorService,

        pipeline,

        gaiaOptions,

        serverConfig.getGaiaConfDir(),

        detectClassPathResourcesToStage(GaiaJobInvoker.class.getClassLoader()));

  }



  static JobInvocation createJobInvocation(

      String invocationId,

      String retrievalToken,

      ListeningExecutorService executorService,

      RunnerApi.Pipeline pipeline,

      GaiaPipelineOptions gaiaOptions,

      @Nullable String confDir,

      List<String> filesToStage) {

    JobInfo jobInfo =

        JobInfo.create(

            invocationId,

            gaiaOptions.getJobName(),

            retrievalToken,

            PipelineOptionsTranslation.toProto(gaiaOptions));

    GaiaPipelineRunner pipelineRunner =

        new GaiaPipelineRunner(gaiaOptions, confDir, filesToStage);

    return new JobInvocation(jobInfo, executorService, pipeline, pipelineRunner);

  }

}